from typing import Union, Optional
from pydantic import BaseModel, Field

class Joke(BaseModel):
    """笑话内容"""
    setup: str = Field(description="笑话的铺垫部分")
    punchline: str = Field(description="笑话的笑点部分")
    rating: Optional[int] = Field(default=None, description="有趣程度，1-10分")

class ConversationalResponse(BaseModel):
    """普通对话回应"""
    response: str = Field(description="对用户查询的对话式回应")

class FinalResponse(BaseModel):
    """最终输出，可以是笑话或普通回应"""
    final_output: Union[Joke, ConversationalResponse]

from langchain_ollama import ChatOllama
llm = ChatOllama(model="qwen3:8b", temperature=0.5, reasoning=False)
structured_llm = llm.with_structured_output(FinalResponse)

# 测试不同类型的输入
print("=" * 30)
print("请求笑话:")
result1 = structured_llm.invoke("给我讲个关于程序员的笑话")
print(result1)

print("=" * 30)
print("普通对话:")
result2 = structured_llm.invoke("你今天怎么样？")
print(result2)